// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.

#ifndef EIGEN_CXX11_TENSOR_TENSOR_BLOCK_H
#define EIGEN_CXX11_TENSOR_TENSOR_BLOCK_H

namespace Eigen {
namespace internal {

// -------------------------------------------------------------------------- //
// Forward declarations for templates defined below.
template<typename Scalar, typename IndexType, int NumDims, int Layout>
class TensorBlockIO;

// -------------------------------------------------------------------------- //
// Helper function to compute strides for densely stored buffer of given
// dimensions.

// TODO(ezhulenev): We compute strides 1000 times in different evaluators, use
// this function instead everywhere.
template<int Layout, typename IndexType, int NumDims>
EIGEN_ALWAYS_INLINE DSizes<IndexType, NumDims>
strides(const DSizes<IndexType, NumDims>& dimensions)
{
	DSizes<IndexType, NumDims> strides;
	if (NumDims == 0)
		return strides;

	// TODO(ezhulenev): Use templates to unroll this loop (similar to
	// h_array_reduce in CXX11meta.h)? Benchmark it.
	if (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
		strides[0] = 1;
		for (int i = 1; i < NumDims; ++i) {
			strides[i] = strides[i - 1] * dimensions[i - 1];
		}
	} else {
		strides[NumDims - 1] = 1;
		for (int i = NumDims - 2; i >= 0; --i) {
			strides[i] = strides[i + 1] * dimensions[i + 1];
		}
	}

	return strides;
}

template<int Layout, typename IndexType, size_t NumDims>
EIGEN_ALWAYS_INLINE DSizes<IndexType, NumDims>
strides(const Eigen::array<IndexType, NumDims>& dimensions)
{
	return strides<Layout>(DSizes<IndexType, NumDims>(dimensions));
}

template<int Layout, std::ptrdiff_t... Indices>
EIGEN_STRONG_INLINE DSizes<std::ptrdiff_t, sizeof...(Indices)>
strides(const Sizes<Indices...>& sizes)
{
	return strides<Layout>(DSizes<std::ptrdiff_t, sizeof...(Indices)>(sizes));
}

// -------------------------------------------------------------------------- //

// Tensor block shape type defines what are the shape preference for the blocks
// extracted from the larger tensor.
//
// Example: blocks of 100 elements from the large 100x100 tensor:
// - tensor: 100x100
// - target_block_size: 100
//
// TensorBlockShapeType:
//  - kUniformAllDims: 100 blocks of size 10x10
//  - kSkewedInnerDims: 100 blocks of size 100x1 (or 1x100 depending on a column
//                      or row major layout)
enum class TensorBlockShapeType
{
	kUniformAllDims,
	kSkewedInnerDims
};

struct TensorBlockResourceRequirements
{
	TensorBlockShapeType shape_type; // target block shape
	size_t size;					 // target block size
	TensorOpCost cost_per_coeff;	 // cost of computing a single block element

#ifdef EIGEN_HIPCC
	// For HIPCC, we need to explicitly declare as a "device fun", the constructor
	// which is implicitly invoked in the "merge" / "any" routines. else HIPCC
	// errors out complaining about the lack of a matching constructor
	EIGEN_DEVICE_FUNC
	TensorBlockResourceRequirements(TensorBlockShapeType shape_type_, size_t size_, TensorOpCost cost_)
		: shape_type(shape_type_)
		, size(size_)
		, cost_per_coeff(cost_)
	{
	}
#endif

	template<typename Scalar>
	EIGEN_DEVICE_FUNC static TensorBlockResourceRequirements withShapeAndSize(TensorBlockShapeType shape_type,
																			  size_t size_in_bytes,
																			  TensorOpCost cost)
	{
		const size_t size = numext::maxi(size_t(1), size_in_bytes / sizeof(Scalar));
		return { shape_type, size, cost };
	}

	template<typename Scalar>
	EIGEN_DEVICE_FUNC static TensorBlockResourceRequirements withShapeAndSize(TensorBlockShapeType shape_type,
																			  size_t size_in_bytes)
	{
		// This default cost per coefficient is valid for most materialized tensor
		// block evaluation implementations, because they typically just read
		// coefficients from the underlying tensor storage, and write to the tensor
		// block buffer (scratch or destination memory, reads and writes have linear
		// access pattern). We ignore the fixed cost of block evaluation, because in
		// practice it should negligible.
		//
		// Lazy block evaluation adds the cost of calling a functor for each
		// coefficient.
		//
		// All non-trivial block evaluation implementations must provide their own
		// cost approximation (e.g. shuffling inner dimension has a much higher cost
		// because it reads memory randomly, although the total number of moved
		// bytes is the same).
		return withShapeAndSize<Scalar>(shape_type,
										size_in_bytes,
										{ /*bytes_loaded=*/sizeof(Scalar),
										  /*bytes_stored=*/sizeof(Scalar),
										  /*compute_cycles=*/0 });
	}

	template<typename Scalar>
	EIGEN_DEVICE_FUNC static TensorBlockResourceRequirements skewed(size_t size_in_bytes)
	{
		return withShapeAndSize<Scalar>(TensorBlockShapeType::kSkewedInnerDims, size_in_bytes);
	}

	template<typename Scalar>
	EIGEN_DEVICE_FUNC static TensorBlockResourceRequirements uniform(size_t size_in_bytes)
	{
		return withShapeAndSize<Scalar>(TensorBlockShapeType::kUniformAllDims, size_in_bytes);
	}

	EIGEN_DEVICE_FUNC
	static EIGEN_STRONG_INLINE TensorBlockResourceRequirements merge(const TensorBlockResourceRequirements& lhs,
																	 const TensorBlockResourceRequirements& rhs)
	{
		return { merge(lhs.shape_type, rhs.shape_type),			  // shape_type
				 merge(lhs.size, rhs.size),						  // size
				 merge(lhs.cost_per_coeff, rhs.cost_per_coeff) }; // cost_per_coeff
	}

	EIGEN_DEVICE_FUNC TensorBlockResourceRequirements& addCostPerCoeff(TensorOpCost cost)
	{
		cost_per_coeff += cost;
		return *this;
	}

	// This is a resource requirement that should be returned from expressions
	// that do not have any block evaluation preference (e.g. default tensor
	// expression with raw buffer access).
	EIGEN_DEVICE_FUNC
	static EIGEN_STRONG_INLINE TensorBlockResourceRequirements any()
	{
		return { TensorBlockShapeType::kUniformAllDims, 1, { 0, 0, 0 } };
	}

  private:
	using Requirements = TensorBlockResourceRequirements;

	EIGEN_DEVICE_FUNC
	static EIGEN_STRONG_INLINE size_t merge(size_t lhs_size, size_t rhs_size)
	{
		return numext::maxi(lhs_size, rhs_size);
	}

	EIGEN_DEVICE_FUNC
	static EIGEN_STRONG_INLINE TensorBlockShapeType merge(TensorBlockShapeType lhs, TensorBlockShapeType rhs)
	{
		return (lhs == TensorBlockShapeType::kSkewedInnerDims || rhs == TensorBlockShapeType::kSkewedInnerDims)
				   ? TensorBlockShapeType::kSkewedInnerDims
				   : TensorBlockShapeType::kUniformAllDims;
	}

	EIGEN_DEVICE_FUNC
	static EIGEN_STRONG_INLINE TensorOpCost merge(TensorOpCost lhs_cost, TensorOpCost rhs_cost)
	{
		return lhs_cost + rhs_cost;
	}
};

// -------------------------------------------------------------------------- //
// TensorBlockDescriptor specifies a block offset within a tensor and the block
// sizes along each of the tensor dimensions.

template<int NumDims, typename IndexType = Eigen::Index>
class TensorBlockDescriptor
{
  public:
	typedef DSizes<IndexType, NumDims> Dimensions;

	// If we evaluate a Tensor assignment, and expression on the left, already has
	// a memory buffer, then we might do performance optimization, and evaluate
	// the root expression directly into the final output memory. Some time it's
	// possible to reuse it for materializing subexpressions inside an expression
	// tree, to to avoid dynamic memory allocation.
	//
	// The pointer type of the underlying storage is erased, because passing
	// Scalar type through all the expression evaluation layers is way too many
	// templates. In practice destination buffer type should always match the
	// evaluated expression scalar type.
	class DestinationBuffer
	{
	  public:
		enum DestinationBufferKind : int
		{
			// The above explicit specification of "int" as the enum basetype is
			// needed to get around a HIPCC link error ("the field type is not
			// amp-compatible")
			// which is issued for class members with the enum type.
			// TODO(rocm):
			// remove the "int" basetype once HIPCC has been fixed to not error out
			// in the above scenario.

			// Destination buffer is not defined (`m_data` == nullptr).
			kEmpty,

			// Tensor block defined by an owning tensor block descriptor can fit
			// contiguously into the destination buffer. In this case it's safe to
			// materialize tensor block in the destination buffer, wrap it in a
			// TensorMap, and use to build Eigen expression on top of it.
			kContiguous,

			// Destination buffer strides do not match strides of the contiguously
			// stored block, and it's impossible to define a TensorMap over this
			// buffer. However if we are evaluating a root of an expression tree, we
			// still can materialize an output into this destination, because we can
			// guarantee that no one will ever access it through block API.
			//
			// In theory it is possible to build valid TensorStriding<TensorMap>
			// expression on top of this destination buffer, however it has
			// inefficient coeff/packet access, and defeats the purpose of fast block
			// evaluation API.
			kStrided
		};

		template<typename Scalar>
		Scalar* data() const
		{
			eigen_assert(m_data_type_size == sizeof(Scalar));
			return static_cast<Scalar*>(m_data);
		}

		const Dimensions& strides() const { return m_strides; }
		const DestinationBufferKind& kind() const { return m_kind; }

	  private:
		friend class TensorBlockDescriptor;

		DestinationBuffer()
			: m_data(NULL)
			, m_data_type_size(0)
			, m_kind(kEmpty)
		{
		}

		template<typename Scalar>
		DestinationBuffer(Scalar* data, const Dimensions& strides, DestinationBufferKind kind)
			: m_data(static_cast<void*>(data))
			, m_data_type_size(sizeof(Scalar))
			, m_strides(strides)
			, m_kind(kind)
		{
		}

		template<int Layout, typename Scalar>
		static DestinationBuffer make(const TensorBlockDescriptor& desc, Scalar* data, const Dimensions& strides)
		{
			return DestinationBuffer(data, strides, kind<Layout>(desc, strides));
		}

		template<int Layout>
		static DestinationBufferKind kind(const TensorBlockDescriptor& desc, const Dimensions& strides)
		{
			const Dimensions& desc_dims = desc.dimensions();
			const Dimensions& desc_strides = internal::strides<Layout>(desc_dims);
			for (int i = 0; i < NumDims; ++i) {
				if (desc_dims[i] == 1)
					continue;
				if (desc_strides[i] != strides[i])
					return kStrided;
			}
			return kContiguous;
		}

		// Storage pointer is type erased, to reduce template bloat, but we still
		// keep the size of the underlying element type for error checking.
		void* m_data;
		size_t m_data_type_size;

		// Destination buffer dimensions always match the dimensions of a tensor
		// block descriptor it belongs to, however strides might be different.
		Dimensions m_strides;

		DestinationBufferKind m_kind;
	};

	TensorBlockDescriptor(const IndexType offset, const Dimensions& dimensions, const DestinationBuffer& destination)
		: m_offset(offset)
		, m_dimensions(dimensions)
		, m_destination(destination)
	{
	}

	TensorBlockDescriptor(const IndexType offset, const Dimensions& dimensions)
		: m_offset(offset)
		, m_dimensions(dimensions)
		, m_destination(DestinationBuffer())
	{
	}

	IndexType offset() const { return m_offset; }
	const Dimensions& dimensions() const { return m_dimensions; }
	IndexType dimension(int index) const { return m_dimensions[index]; }
	IndexType size() const { return array_prod<IndexType>(m_dimensions); }

	const DestinationBuffer& destination() const { return m_destination; }

	template<int Layout, typename Scalar>
	void AddDestinationBuffer(Scalar* dst_base, const Dimensions& dst_strides)
	{
		eigen_assert(dst_base != NULL);
		m_destination = DestinationBuffer::template make<Layout>(*this, dst_base, dst_strides);
	}

	template<int Layout, typename Scalar, typename DstStridesIndexType>
	void AddDestinationBuffer(Scalar* dst_base, const DSizes<DstStridesIndexType, NumDims>& dst_strides)
	{
		// DSizes constructor will do index type promotion if it's safe.
		AddDestinationBuffer<Layout>(dst_base, Dimensions(dst_strides));
	}

	TensorBlockDescriptor& DropDestinationBuffer()
	{
		m_destination.m_data = NULL;
		m_destination.m_kind = DestinationBuffer::kEmpty;
		return *this;
	}

	bool HasDestinationBuffer() const { return m_destination.kind() != DestinationBuffer::kEmpty; }

	// Returns a copy of `*this` with updated offset.
	TensorBlockDescriptor WithOffset(IndexType offset) const
	{
		return TensorBlockDescriptor(offset, m_dimensions, m_destination);
	}

  private:
	// Offset and dimensions are immutable after construction. Block descriptor
	// can only be mutated by adding or dropping destination.
	const IndexType m_offset;
	const Dimensions m_dimensions;
	DestinationBuffer m_destination;
};

// -------------------------------------------------------------------------- //
// TensorBlockMapper is responsible for iterating over the blocks of a tensor.

template<int NumDims, int Layout, typename IndexType = Eigen::Index>
class TensorBlockMapper
{
	typedef TensorBlockDescriptor<NumDims, IndexType> BlockDescriptor;

  public:
	typedef DSizes<IndexType, NumDims> Dimensions;

	TensorBlockMapper() = default;
	TensorBlockMapper(const DSizes<IndexType, NumDims>& dimensions, const TensorBlockResourceRequirements& requirements)
		: m_tensor_dimensions(dimensions)
		, m_requirements(requirements)
	{
		// Compute block dimensions and the total number of blocks.
		InitializeBlockDimensions();
	}

	EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE IndexType blockCount() const { return m_total_block_count; }

	EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE IndexType blockTotalSize() const { return m_block_dimensions.TotalSize(); }

	EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const DSizes<IndexType, NumDims>& blockDimensions() const
	{
		return m_block_dimensions;
	}

	EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE BlockDescriptor blockDescriptor(IndexType block_index) const
	{
		static const bool isColMajor = Layout == static_cast<int>(ColMajor);

		IndexType offset = 0;
		DSizes<IndexType, NumDims> dimensions;

		if (NumDims == 0)
			return BlockDescriptor(offset, dimensions);

		// Iterate outer -> inner dimensions.
		for (int i = NumDims - 1; i >= 0; --i) {
			const int dim = isColMajor ? i : NumDims - i - 1;

			const IndexType idx = block_index / m_block_strides[dim];
			block_index -= idx * m_block_strides[dim];

			const IndexType coord = idx * m_block_dimensions[dim];
			dimensions[dim] = numext::mini(m_tensor_dimensions[dim] - coord, m_block_dimensions[dim]);
			offset += coord * m_tensor_strides[dim];
		}

		return { offset, dimensions };
	}

  private:
	void InitializeBlockDimensions()
	{
		// Requested block shape and size.
		const TensorBlockShapeType shape_type = m_requirements.shape_type;
		IndexType target_block_size = numext::maxi<IndexType>(1, static_cast<IndexType>(m_requirements.size));

		IndexType tensor_size = m_tensor_dimensions.TotalSize();

		// Corner case: one of the dimensions is zero. Logic below is too complex
		// to handle this case on a general basis, just use unit block size.
		// Note: we must not yield blocks with zero dimensions (recipe for
		// overflows/underflows, divisions by zero and NaNs later).
		if (tensor_size == 0) {
			for (int i = 0; i < NumDims; ++i) {
				m_block_dimensions[i] = 1;
			}
			m_total_block_count = 0;
			return;
		}

		// If tensor fits into a target block size, evaluate it as a single block.
		if (tensor_size <= target_block_size) {
			m_block_dimensions = m_tensor_dimensions;
			m_total_block_count = 1;
			// The only valid block index is `0`, and in this case we do not need
			// to compute real strides for tensor or blocks (see blockDescriptor).
			for (int i = 0; i < NumDims; ++i) {
				m_tensor_strides[i] = 0;
				m_block_strides[i] = 1;
			}
			return;
		}

		static const bool isColMajor = Layout == static_cast<int>(ColMajor);

		// Block shape skewed towards inner dimension.
		if (shape_type == TensorBlockShapeType::kSkewedInnerDims) {
			IndexType coeff_to_allocate = target_block_size;

			for (int i = 0; i < NumDims; ++i) {
				const int dim = isColMajor ? i : NumDims - i - 1;
				m_block_dimensions[dim] = numext::mini(coeff_to_allocate, m_tensor_dimensions[dim]);
				coeff_to_allocate =
					divup(coeff_to_allocate, numext::maxi(static_cast<IndexType>(1), m_block_dimensions[dim]));
			}
			eigen_assert(coeff_to_allocate == 1);

		} else if (shape_type == TensorBlockShapeType::kUniformAllDims) {
			// Tensor will not fit within 'target_block_size' budget: calculate tensor
			// block dimension sizes based on "square" dimension size target.
			const IndexType dim_size_target = convert_index<IndexType>(
				std::pow(static_cast<float>(target_block_size), 1.0f / static_cast<float>(m_block_dimensions.rank())));

			for (int i = 0; i < NumDims; ++i) {
				// TODO(andydavis) Adjust the inner most 'block_dim_size' to make it
				// a multiple of the packet size. Note that reducing
				// 'block_dim_size' in this manner can increase the number of
				// blocks, and so will amplify any per-block overhead.
				m_block_dimensions[i] = numext::mini(dim_size_target, m_tensor_dimensions[i]);
			}

			// Add any un-allocated coefficients to inner dimension(s).
			IndexType total_size = m_block_dimensions.TotalSize();
			for (int i = 0; i < NumDims; ++i) {
				const int dim = isColMajor ? i : NumDims - i - 1;

				if (m_block_dimensions[dim] < m_tensor_dimensions[dim]) {
					const IndexType total_size_other_dims = total_size / m_block_dimensions[dim];
					const IndexType alloc_avail = divup<IndexType>(target_block_size, total_size_other_dims);
					if (alloc_avail == m_block_dimensions[dim]) {
						// Insufficient excess coefficients to allocate.
						break;
					}
					m_block_dimensions[dim] = numext::mini(m_tensor_dimensions[dim], alloc_avail);
					total_size = total_size_other_dims * m_block_dimensions[dim];
				}
			}

		} else {
			eigen_assert(false); // unknown block shape
		}

		eigen_assert(m_block_dimensions.TotalSize() >=
					 numext::mini<IndexType>(target_block_size, m_tensor_dimensions.TotalSize()));

		// Calculate block counts by dimension and total block count.
		DSizes<IndexType, NumDims> block_count;
		for (int i = 0; i < NumDims; ++i) {
			block_count[i] = divup(m_tensor_dimensions[i], m_block_dimensions[i]);
		}
		m_total_block_count = array_prod(block_count);

		// Calculate block strides (used for enumerating blocks).
		m_tensor_strides = strides<Layout>(m_tensor_dimensions);
		m_block_strides = strides<Layout>(block_count);
	}

	DSizes<IndexType, NumDims> m_tensor_dimensions;
	TensorBlockResourceRequirements m_requirements;

	DSizes<IndexType, NumDims> m_block_dimensions;
	IndexType m_total_block_count;

	DSizes<IndexType, NumDims> m_tensor_strides;
	DSizes<IndexType, NumDims> m_block_strides;
};

// -------------------------------------------------------------------------- //
// TensorBlockScratchAllocator is responsible for allocating temporary buffers
// for block evaluation (output or input block materialization). Given that
// Eigen expression traversal order is deterministic, all temporary allocations
// are happening in the same order, and usually have exactly the same size.
// Scratch allocator keeps a trace of all dynamic allocations, and after the
// first block evaluation is completed, we should be able to reuse all the
// temporary buffers for the next block evaluation.

template<typename Device>
class TensorBlockScratchAllocator
{
  public:
	explicit TensorBlockScratchAllocator(const Device& device)
		: m_device(device)
		, m_allocation_index(0)
	{
	}

	~TensorBlockScratchAllocator()
	{
		for (size_t i = 0; i < m_allocations.size(); ++i) {
			m_device.deallocate(m_allocations[i].ptr);
		}
	}

	void* allocate(size_t size)
	{
		// TODO(ezhulenev): Remove when replaced with inlined vector.
		if (m_allocations.capacity() == 0)
			m_allocations.reserve(8);

		// Check if we already have an existing allocation att current index.
		const int num_allocations = static_cast<int>(m_allocations.size());
		const bool has_allocation = m_allocation_index < num_allocations;

		// Allocation index can't be larger than the number of allocations.
		eigen_assert(m_allocation_index <= num_allocations);

		// If we have existing allocation, and its size is larger or equal to
		// requested size, we do nothing.

		// If current allocation can't fit requested size, we deallocate it, and
		// replace with a larger allocation.
		if (has_allocation && m_allocations[m_allocation_index].size < size) {
			m_device.deallocate(m_allocations[m_allocation_index].ptr);
			m_allocations[m_allocation_index].ptr = m_device.allocate(size);
			m_allocations[m_allocation_index].size = size;
		}

		// Make a new allocation if we don't have and existing one.
		if (!has_allocation) {
			Allocation allocation;
			allocation.ptr = m_device.allocate(size);
			allocation.size = size;
			m_allocations.push_back(allocation);
		}

		eigen_assert(m_allocations[m_allocation_index].ptr != NULL);
		eigen_assert(m_allocations[m_allocation_index].size >= size);

		return m_allocations[m_allocation_index++].ptr;
	}

	void reset() { m_allocation_index = 0; }

  private:
	struct Allocation
	{
		void* ptr;
		size_t size;
	};

	const Device& m_device;
	int m_allocation_index;
	// TODO(ezhulenev): This should be an inlined vector.
	std::vector<Allocation> m_allocations;
};

// -------------------------------------------------------------------------- //
// TensorBlockKind represents all possible block kinds, that can be produced by
// TensorEvaluator::evalBlock function.
enum TensorBlockKind
{
	// Tensor block that is a lazy expression that must be assigned to a
	// destination using TensorBlockAssign.
	kExpr,

	// Tensor block that is a view into a memory buffer owned by an underlying
	// Tensor expression (e.g. it can be a view into a Tensor buffer).
	kView,

	// Tensor block that was materialized in a scratch memory buffer, allocated
	// with TensorBlockScratchAllocator. This block must be copied to a
	// destination, similar to a block of `kExpr` type.
	kMaterializedInScratch,

	// Tensor block that was materialized directly into the final output memory
	// buffer. For example if the left side of an assignment is a Tensor, we can
	// directly materialize the block in the destination memory.
	//
	// If strides in the output buffer do not match tensor block strides, the
	// Tensor expression will be invalid, and should not be used by
	// TensorBlockAssign or for constructing another block expression.
	kMaterializedInOutput
};

// -------------------------------------------------------------------------- //
// TensorBlockNotImplemented should be used to defined TensorBlock typedef in
// TensorEvaluators that do not support block evaluation.

class TensorBlockNotImplemented
{
  public:
	typedef void XprType;
};

// -------------------------------------------------------------------------- //
// XprScalar extracts Scalar type from the Eigen expressions (if expression type
// is not void). It's required to be able to define lazy block expression for
// argument types, that do not support block evaluation.

template<typename XprType>
struct XprScalar
{
	typedef typename XprType::Scalar type;
};
template<>
struct XprScalar<void>
{
	typedef void type;
};

// -------------------------------------------------------------------------- //
// TensorMaterializedBlock is a fully evaluated block of the original tensor,
// and XprType is just a TensorMap over the data. This block type is typically
// used to materialize blocks of tensor expressions, that can't be efficiently
// represented as lazy Tensor expressions with fast coeff/packet operations,
// e.g. we materialize all broadcasts into evaluated blocks.
//
// TensorMaterializedBlock does not own its memory buffer, it's either a memory
// buffer that backs the original expression (e.g. block is just a view into a
// Tensor), or a memory buffer allocated with scratch allocator, and in this
// case the scratch allocator will deallocate it at the end of block based
// expression execution.
//
// If the block was evaluated directly into the output buffer, and strides in
// the output buffer do not match block strides, the TensorMap expression will
// be invalid, and should never be used in block assignment or any other tensor
// expression.

template<typename Scalar, int NumDims, int Layout, typename IndexType = Eigen::Index>
class TensorMaterializedBlock
{
  public:
	typedef DSizes<IndexType, NumDims> Dimensions;
	typedef TensorMap<const Tensor<Scalar, NumDims, Layout>> XprType;

	TensorMaterializedBlock(TensorBlockKind kind,
							const Scalar* data,
							const Dimensions& dimensions,
							bool valid_expr = true)
		: m_kind(kind)
		, m_data(data)
		, m_dimensions(dimensions)
		, m_expr(m_data, m_dimensions)
		, m_valid_expr(valid_expr)
	{
		eigen_assert(m_kind == internal::TensorBlockKind::kView ||
					 m_kind == internal::TensorBlockKind::kMaterializedInScratch ||
					 m_kind == internal::TensorBlockKind::kMaterializedInOutput);
	}

	TensorBlockKind kind() const { return m_kind; }
	// NOTE(ezhulenev): Returning XprType by value like in other block types
	// causes asan failures. The theory is that XprType::Nested doesn't work
	// properly for TensorMap.
	const XprType& expr() const
	{
		eigen_assert(m_valid_expr);
		return m_expr;
	}
	const Scalar* data() const { return m_data; }
	void cleanup() {}

	typedef internal::TensorBlockDescriptor<NumDims, IndexType> TensorBlockDesc;

	// TensorMaterializedBlock can be backed by different types of storage:
	//
	//   (1) Contiguous block of memory allocated with scratch allocator.
	//   (2) Contiguous block of memory reused from tensor block descriptor
	//       destination buffer.
	//   (3) Strided block of memory reused from tensor block descriptor
	//       destination buffer.
	//
	class Storage
	{
	  public:
		Scalar* data() const { return m_data; }
		const Dimensions& dimensions() const { return m_dimensions; }
		const Dimensions& strides() const { return m_strides; }

		TensorMaterializedBlock AsTensorMaterializedBlock() const
		{
			return TensorMaterializedBlock(m_materialized_in_output ? internal::TensorBlockKind::kMaterializedInOutput
																	: internal::TensorBlockKind::kMaterializedInScratch,
										   m_data,
										   m_dimensions,
										   !m_strided_storage);
		}

	  private:
		friend class TensorMaterializedBlock;

		Storage(Scalar* data,
				const Dimensions& dimensions,
				const Dimensions& strides,
				bool materialized_in_output,
				bool strided_storage)
			: m_data(data)
			, m_dimensions(dimensions)
			, m_strides(strides)
			, m_materialized_in_output(materialized_in_output)
			, m_strided_storage(strided_storage)
		{
		}

		Scalar* m_data;
		Dimensions m_dimensions;
		Dimensions m_strides;
		bool m_materialized_in_output;
		bool m_strided_storage;
	};

	// Creates a storage for materialized block either from the block descriptor
	// destination buffer, or allocates a new buffer with scratch allocator.
	template<typename TensorBlockScratch>
	EIGEN_STRONG_INLINE static Storage prepareStorage(TensorBlockDesc& desc,
													  TensorBlockScratch& scratch,
													  bool allow_strided_storage = false)
	{
		// Try to reuse destination as an output block buffer.
		typedef typename TensorBlockDesc::DestinationBuffer DestinationBuffer;

		if (desc.destination().kind() == DestinationBuffer::kContiguous) {
			Scalar* buffer = desc.destination().template data<Scalar>();
			desc.DropDestinationBuffer();
			return Storage(buffer,
						   desc.dimensions(),
						   internal::strides<Layout>(desc.dimensions()),
						   /*materialized_in_output=*/true,
						   /*strided_storage=*/false);

		} else if (desc.destination().kind() == DestinationBuffer::kStrided && allow_strided_storage) {
			Scalar* buffer = desc.destination().template data<Scalar>();
			desc.DropDestinationBuffer();
			return Storage(buffer,
						   desc.dimensions(),
						   desc.destination().strides(),
						   /*materialized_in_output=*/true,
						   /*strided_storage=*/true);

		} else {
			void* mem = scratch.allocate(desc.size() * sizeof(Scalar));
			return Storage(static_cast<Scalar*>(mem),
						   desc.dimensions(),
						   internal::strides<Layout>(desc.dimensions()),
						   /*materialized_in_output=*/false,
						   /*strided_storage=*/false);
		}
	}

	// Creates a materialized block for the given descriptor from a memory buffer.
	template<typename DataDimensions, typename TensorBlockScratch>
	EIGEN_STRONG_INLINE static TensorMaterializedBlock materialize(const Scalar* data,
																   const DataDimensions& data_dims,
																   TensorBlockDesc& desc,
																   TensorBlockScratch& scratch)
	{
		eigen_assert(array_size<DataDimensions>::value == desc.dimensions().size());

		// If a tensor block dimensions covers a contiguous block of the underlying
		// memory, we can skip block buffer memory allocation, and construct a block
		// from existing `data` memory buffer.
		//
		// Example: (RowMajor layout)
		//   data_dims:          [11, 12, 13, 14]
		//   desc.dimensions():  [1,   1,  3, 14]
		//
		// In this case we can construct a TensorBlock starting at
		// `data + desc.offset()`, with a `desc.dimensions()` block sizes.
		static const bool is_col_major = Layout == ColMajor;

		// Find out how many inner dimensions have a matching size.
		int num_matching_inner_dims = 0;
		for (int i = 0; i < NumDims; ++i) {
			int dim = is_col_major ? i : NumDims - i - 1;
			if (data_dims[dim] != desc.dimensions()[dim])
				break;
			++num_matching_inner_dims;
		}

		// All the outer dimensions must be of size `1`, except a single dimension
		// before the matching inner dimension (`3` in the example above).
		bool can_use_direct_access = true;
		for (int i = num_matching_inner_dims + 1; i < NumDims; ++i) {
			int dim = is_col_major ? i : NumDims - i - 1;
			if (desc.dimension(dim) != 1) {
				can_use_direct_access = false;
				break;
			}
		}

		if (can_use_direct_access) {
			const Scalar* block_start = data + desc.offset();
			return TensorMaterializedBlock(internal::TensorBlockKind::kView, block_start, desc.dimensions());

		} else {
			// Reuse destination buffer or allocate new buffer with scratch allocator.
			const Storage storage = prepareStorage(desc, scratch);

			typedef internal::TensorBlockIO<Scalar, IndexType, NumDims, Layout> TensorBlockIO;
			typedef typename TensorBlockIO::Dst TensorBlockIODst;
			typedef typename TensorBlockIO::Src TensorBlockIOSrc;

			TensorBlockIOSrc src(internal::strides<Layout>(Dimensions(data_dims)), data, desc.offset());
			TensorBlockIODst dst(storage.dimensions(), storage.strides(), storage.data());

			TensorBlockIO::Copy(dst, src);
			return storage.AsTensorMaterializedBlock();
		}
	}

  private:
	TensorBlockKind m_kind;
	const Scalar* m_data;
	Dimensions m_dimensions;
	XprType m_expr;
	bool m_valid_expr;
};

// -------------------------------------------------------------------------- //
// TensorCwiseUnaryBlock is a lazy tensor expression block that applies UnaryOp
// functor to the blocks produced by the underlying Tensor expression.

template<typename UnaryOp, typename ArgTensorBlock>
class TensorCwiseUnaryBlock
{
	static const bool NoArgBlockAccess = internal::is_void<typename ArgTensorBlock::XprType>::value;

  public:
	typedef typename conditional<NoArgBlockAccess,
								 void,
								 TensorCwiseUnaryOp<UnaryOp, const typename ArgTensorBlock::XprType>>::type XprType;

	typedef typename XprScalar<XprType>::type Scalar;

	TensorCwiseUnaryBlock(const ArgTensorBlock& arg_block, const UnaryOp& functor)
		: m_arg_block(arg_block)
		, m_functor(functor)
	{
	}

	TensorBlockKind kind() const { return internal::TensorBlockKind::kExpr; }

	XprType expr() const { return XprType(m_arg_block.expr(), m_functor); }
	const Scalar* data() const { return NULL; }
	void cleanup() { m_arg_block.cleanup(); }

  private:
	ArgTensorBlock m_arg_block;
	UnaryOp m_functor;
};

// -------------------------------------------------------------------------- //
// TensorCwiseUnaryBlock is a lazy tensor expression block that applies BinaryOp
// functor to the blocks produced by the underlying Tensor expression.

template<typename BinaryOp, typename LhsTensorBlock, typename RhsTensorBlock>
class TensorCwiseBinaryBlock
{
	static const bool NoArgBlockAccess = internal::is_void<typename LhsTensorBlock::XprType>::value ||
										 internal::is_void<typename RhsTensorBlock::XprType>::value;

  public:
	typedef typename conditional<
		NoArgBlockAccess,
		void,
		TensorCwiseBinaryOp<BinaryOp, const typename LhsTensorBlock::XprType, const typename RhsTensorBlock::XprType>>::
		type XprType;

	typedef typename XprScalar<XprType>::type Scalar;

	TensorCwiseBinaryBlock(const LhsTensorBlock& left_block, const RhsTensorBlock& right_block, const BinaryOp& functor)
		: m_left_block(left_block)
		, m_right_block(right_block)
		, m_functor(functor)
	{
	}

	TensorBlockKind kind() const { return internal::TensorBlockKind::kExpr; }

	XprType expr() const { return XprType(m_left_block.expr(), m_right_block.expr(), m_functor); }

	const Scalar* data() const { return NULL; }

	void cleanup()
	{
		m_left_block.cleanup();
		m_right_block.cleanup();
	}

  private:
	LhsTensorBlock m_left_block;
	RhsTensorBlock m_right_block;
	BinaryOp m_functor;
};

// -------------------------------------------------------------------------- //
// TensorUnaryExprBlock is a lazy tensor expression block that can construct
// an arbitrary tensor expression from a block of the underlying type (this is a
// generalization of the TensorCwiseUnaryBlock for arbitrary expressions).

template<typename BlockFactory, typename ArgTensorBlock>
class TensorUnaryExprBlock
{
	typedef typename ArgTensorBlock::XprType ArgXprType;
	static const bool NoArgBlockAccess = internal::is_void<ArgXprType>::value;

  public:
	typedef
		typename conditional<NoArgBlockAccess, void, typename BlockFactory::template XprType<ArgXprType>::type>::type
			XprType;

	typedef typename XprScalar<XprType>::type Scalar;

	TensorUnaryExprBlock(const ArgTensorBlock& arg_block, const BlockFactory& factory)
		: m_arg_block(arg_block)
		, m_factory(factory)
	{
	}

	TensorBlockKind kind() const { return internal::TensorBlockKind::kExpr; }
	XprType expr() const { return m_factory.expr(m_arg_block.expr()); }
	const Scalar* data() const { return NULL; }
	void cleanup() { m_arg_block.cleanup(); }

  private:
	ArgTensorBlock m_arg_block;
	BlockFactory m_factory;
};

// -------------------------------------------------------------------------- //
// TensorTernaryExprBlock is a lazy tensor expression block that can construct
// an arbitrary tensor expression from three blocks of the underlying type.

template<typename BlockFactory, typename Arg1TensorBlock, typename Arg2TensorBlock, typename Arg3TensorBlock>
class TensorTernaryExprBlock
{
	typedef typename Arg1TensorBlock::XprType Arg1XprType;
	typedef typename Arg2TensorBlock::XprType Arg2XprType;
	typedef typename Arg3TensorBlock::XprType Arg3XprType;

	static const bool NoArgBlockAccess = internal::is_void<Arg1XprType>::value ||
										 internal::is_void<Arg2XprType>::value || internal::is_void<Arg3XprType>::value;

  public:
	typedef
		typename conditional<NoArgBlockAccess,
							 void,
							 typename BlockFactory::template XprType<Arg1XprType, Arg2XprType, Arg3XprType>::type>::type
			XprType;

	typedef typename XprScalar<XprType>::type Scalar;

	TensorTernaryExprBlock(const Arg1TensorBlock& arg1_block,
						   const Arg2TensorBlock& arg2_block,
						   const Arg3TensorBlock& arg3_block,
						   const BlockFactory& factory)
		: m_arg1_block(arg1_block)
		, m_arg2_block(arg2_block)
		, m_arg3_block(arg3_block)
		, m_factory(factory)
	{
	}

	TensorBlockKind kind() const { return internal::TensorBlockKind::kExpr; }
	XprType expr() const { return m_factory.expr(m_arg1_block.expr(), m_arg2_block.expr(), m_arg3_block.expr()); }
	const Scalar* data() const { return NULL; }
	void cleanup()
	{
		m_arg1_block.cleanup();
		m_arg2_block.cleanup();
		m_arg3_block.cleanup();
	}

  private:
	Arg1TensorBlock m_arg1_block;
	Arg2TensorBlock m_arg2_block;
	Arg3TensorBlock m_arg3_block;
	BlockFactory m_factory;
};

// -------------------------------------------------------------------------- //
// StridedLinearBufferCopy provides a method to copy data between two linear
// buffers with different strides, with optimized paths for scatter/gather.

template<typename Scalar, typename IndexType>
class StridedLinearBufferCopy
{
	typedef typename packet_traits<Scalar>::type Packet;
	enum
	{
		Vectorizable = packet_traits<Scalar>::Vectorizable,
		PacketSize = packet_traits<Scalar>::size
	};

  public:
	// Specifying linear copy kind statically gives ~30% speedup for small sizes.
	enum class Kind
	{
		Linear = 0,		 // src_stride == 1 && dst_stride == 1
		Scatter = 1,	 // src_stride == 1 && dst_stride != 1
		FillLinear = 2,	 // src_stride == 0 && dst_stride == 1
		FillScatter = 3, // src_stride == 0 && dst_stride != 1
		Gather = 4,		 // dst_stride == 1
		Random = 5		 // everything else
	};

	struct Dst
	{
		Dst(IndexType o, IndexType s, Scalar* d)
			: offset(o)
			, stride(s)
			, data(d)
		{
		}

		IndexType offset;
		IndexType stride;
		Scalar* data;
	};

	struct Src
	{
		Src(IndexType o, IndexType s, const Scalar* d)
			: offset(o)
			, stride(s)
			, data(d)
		{
		}

		IndexType offset;
		IndexType stride;
		const Scalar* data;
	};

	template<typename StridedLinearBufferCopy::Kind kind>
	static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void Run(const Dst& dst, const Src& src, const size_t count)
	{
		Run<kind>(count, dst.offset, dst.stride, dst.data, src.offset, src.stride, src.data);
	}

  private:
	template<typename StridedLinearBufferCopy::Kind kind>
	static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void Run(const IndexType count,
														  const IndexType dst_offset,
														  const IndexType dst_stride,
														  Scalar* EIGEN_RESTRICT dst_data,
														  const IndexType src_offset,
														  const IndexType src_stride,
														  const Scalar* EIGEN_RESTRICT src_data)
	{
		const Scalar* src = &src_data[src_offset];
		Scalar* dst = &dst_data[dst_offset];

		if (!Vectorizable) {
			for (Index i = 0; i < count; ++i) {
				dst[i * dst_stride] = src[i * src_stride];
			}
			return;
		}

		const IndexType vectorized_size = count - PacketSize;
		IndexType i = 0;

		if (kind == StridedLinearBufferCopy::Kind::Linear) {
			// ******************************************************************** //
			// Linear copy from `src` to `dst`.
			const IndexType unrolled_size = count - 4 * PacketSize;
			eigen_assert(src_stride == 1 && dst_stride == 1);
			for (; i <= unrolled_size; i += 4 * PacketSize) {
				for (int j = 0; j < 4; ++j) {
					Packet p = ploadu<Packet>(src + i + j * PacketSize);
					pstoreu<Scalar, Packet>(dst + i + j * PacketSize, p);
				}
			}
			for (; i <= vectorized_size; i += PacketSize) {
				Packet p = ploadu<Packet>(src + i);
				pstoreu<Scalar, Packet>(dst + i, p);
			}
			for (; i < count; ++i) {
				dst[i] = src[i];
			}
			// ******************************************************************** //
		} else if (kind == StridedLinearBufferCopy::Kind::Scatter) {
			// Scatter from `src` to `dst`.
			eigen_assert(src_stride == 1 && dst_stride != 1);
			for (; i <= vectorized_size; i += PacketSize) {
				Packet p = ploadu<Packet>(src + i);
				pscatter<Scalar, Packet>(dst + i * dst_stride, p, dst_stride);
			}
			for (; i < count; ++i) {
				dst[i * dst_stride] = src[i];
			}
			// ******************************************************************** //
		} else if (kind == StridedLinearBufferCopy::Kind::FillLinear) {
			// Fill `dst` with value at `*src`.
			eigen_assert(src_stride == 0 && dst_stride == 1);
			const IndexType unrolled_size = count - 4 * PacketSize;
			Packet p = pload1<Packet>(src);
			for (; i <= unrolled_size; i += 4 * PacketSize) {
				for (int j = 0; j < 4; ++j) {
					pstoreu<Scalar, Packet>(dst + i + j * PacketSize, p);
				}
			}
			for (; i <= vectorized_size; i += PacketSize) {
				pstoreu<Scalar, Packet>(dst + i, p);
			}
			for (; i < count; ++i) {
				dst[i] = *src;
			}
			// ******************************************************************** //
		} else if (kind == StridedLinearBufferCopy::Kind::FillScatter) {
			// Scatter `*src` into `dst`.
			eigen_assert(src_stride == 0 && dst_stride != 1);
			Packet p = pload1<Packet>(src);
			for (; i <= vectorized_size; i += PacketSize) {
				pscatter<Scalar, Packet>(dst + i * dst_stride, p, dst_stride);
			}
			for (; i < count; ++i) {
				dst[i * dst_stride] = *src;
			}
			// ******************************************************************** //
		} else if (kind == StridedLinearBufferCopy::Kind::Gather) {
			// Gather from `src` into `dst`.
			eigen_assert(dst_stride == 1);
			for (; i <= vectorized_size; i += PacketSize) {
				Packet p = pgather<Scalar, Packet>(src + i * src_stride, src_stride);
				pstoreu<Scalar, Packet>(dst + i, p);
			}
			for (; i < count; ++i) {
				dst[i] = src[i * src_stride];
			}
			// ******************************************************************** //
		} else if (kind == StridedLinearBufferCopy::Kind::Random) {
			// Random.
			for (; i < count; ++i) {
				dst[i * dst_stride] = src[i * src_stride];
			}
		} else {
			eigen_assert(false);
		}
	}
};

// -------------------------------------------------------------------------- //
// TensorBlockIO copies data from `src` tensor block, to the `dst` tensor block.
// It's possible to specify src->dst dimension mapping for the copy operation.
// Dimensions of `dst` specify how many elements have to be copied, for the
// `src` we need to know only stride to navigate through source memory buffer.

template<typename Scalar, typename IndexType, int NumDims, int Layout>
class TensorBlockIO
{
	static const bool IsColMajor = (Layout == ColMajor);

	typedef StridedLinearBufferCopy<Scalar, IndexType> LinCopy;

  public:
	typedef DSizes<IndexType, NumDims> Dimensions;
	typedef DSizes<int, NumDims> DimensionsMap;

	struct Dst
	{
		Dst(const Dimensions& dst_dims, const Dimensions& dst_strides, Scalar* dst, IndexType dst_offset = 0)
			: dims(dst_dims)
			, strides(dst_strides)
			, data(dst)
			, offset(dst_offset)
		{
		}

		Dimensions dims;
		Dimensions strides;
		Scalar* data;
		IndexType offset;
	};

	struct Src
	{
		Src(const Dimensions& src_strides, const Scalar* src, IndexType src_offset = 0)
			: strides(src_strides)
			, data(src)
			, offset(src_offset)
		{
		}

		Dimensions strides;
		const Scalar* data;
		IndexType offset;
	};

	// Copies data to `dst` from `src`, using provided dimensions mapping:
	//
	//   src_dimension_index = dst_to_src_dim_map[dst_dimension_index]
	//
	// Returns the number of copied elements.
	static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE IndexType Copy(const Dst& dst,
																const Src& src,
																const DimensionsMap& dst_to_src_dim_map)
	{
		// Copy single scalar value from `src` to `dst`.
		if (NumDims == 0) {
			*(dst.data + dst.offset) = *(src.data + src.offset);
			return 1;
		}

		// Both `dst` and `src` must have contiguous innermost dimension. We also
		// accept the special case with stride '0', because it's used as a trick to
		// implement broadcasting.
		{
			int inner_dim = IsColMajor ? 0 : NumDims - 1;
			EIGEN_UNUSED_VARIABLE(inner_dim);
			eigen_assert(dst.strides[inner_dim] == 1 || dst.strides[inner_dim] == 0);
			eigen_assert(src.strides[inner_dim] == 1 || src.strides[inner_dim] == 0);
		}

		// Give a shorter name to `dst_to_src_dim_map`.
		const DimensionsMap& dim_map = dst_to_src_dim_map;

		// Do not squeeze reordered inner dimensions.
		int num_squeezable_dims = NumSqueezableInnerDims(dim_map);

		// NOTE: We find the innermost dimension (contiguous in memory) in the dst
		// block, and we write data linearly into that dimension, reading it from
		// the src. If dimensions are reordered, we might end up reading data from
		// the src with `stride != 1`.
		//
		// NOTE: Random-Read/Linear-Write can be up to ~2X faster than
		// Linear-Read/Random-Write: https://stackoverflow.com/a/54935680

		// Find the innermost dimension in the dst whose size is not 1. This is the
		// effective inner dim.
		int num_size_one_inner_dims = 0;
		for (int i = 0; i < num_squeezable_dims; ++i) {
			const int dst_dim = IsColMajor ? i : NumDims - i - 1;
			if (dst.dims[dst_dim] != 1)
				break;
			num_size_one_inner_dims++;
		}

		// If all dimensions are of size 1, just copy a scalar from `src` to `dst`.
		if (num_size_one_inner_dims == NumDims) {
			*(dst.data + dst.offset) = *(src.data + src.offset);
			return 1;
		}

		// Outermost dimension in the dst with `stride == 1` (contiguous in memory).
		const int dst_stride1_dim = IsColMajor ? num_size_one_inner_dims : NumDims - num_size_one_inner_dims - 1;

		// Dimension in the src that corresponds to the dst innermost dimension.
		const int src_dim_for_dst_stride1_dim = NumDims == 0 ? 1 : dim_map[dst_stride1_dim];

		// Size of the innermost dimension (length of contiguous blocks of memory).
		IndexType dst_inner_dim_size = NumDims == 0 ? 1 : dst.dims[dst_stride1_dim];

		// Squeeze multiple inner dims into one if they are contiguous in `dst` and
		// `src` memory, so we can do less linear copy calls.
		for (int i = num_size_one_inner_dims + 1; i < num_squeezable_dims; ++i) {
			const int dst_dim = IsColMajor ? i : NumDims - i - 1;
			const IndexType dst_stride = dst.strides[dst_dim];
			const IndexType src_stride = src.strides[dim_map[dst_dim]];
			if (dst_inner_dim_size == dst_stride && dst_stride == src_stride) {
				dst_inner_dim_size *= dst.dims[dst_dim];
				++num_size_one_inner_dims;
			} else {
				break;
			}
		}

		// Setup strides to read data from `src` and write to `dst`.
		IndexType input_offset = src.offset;
		IndexType output_offset = dst.offset;
		IndexType input_stride = NumDims == 0 ? 1 : src.strides[src_dim_for_dst_stride1_dim];
		IndexType output_stride = NumDims == 0 ? 1 : dst.strides[dst_stride1_dim];

		const int at_least_1_dim = NumDims <= 1 ? 1 : NumDims - 1;
		array<BlockIteratorState, at_least_1_dim> it;

		// Initialize block iterator state. Squeeze away any dimension of size 1.
		int idx = 0; // currently initialized iterator state index
		for (int i = num_size_one_inner_dims; i < NumDims - 1; ++i) {
			const int dst_dim = IsColMajor ? i + 1 : NumDims - i - 2;
			if (dst.dims[dst_dim] == 1)
				continue;

			it[idx].size = dst.dims[dst_dim];
			it[idx].input_stride = src.strides[dim_map[dst_dim]];
			it[idx].output_stride = dst.strides[dst_dim];

			it[idx].input_span = it[idx].input_stride * (it[idx].size - 1);
			it[idx].output_span = it[idx].output_stride * (it[idx].size - 1);

			idx++;
		}

		// Iterate copying data from src to dst.
		const IndexType block_total_size = NumDims == 0 ? 1 : dst.dims.TotalSize();

#define COPY_INNER_DIM(KIND)                                                                                           \
	IndexType num_copied = 0;                                                                                          \
	for (num_copied = 0; num_copied < block_total_size; num_copied += dst_inner_dim_size) {                            \
		LinCopy::template Run<KIND>(typename LinCopy::Dst(output_offset, output_stride, dst.data),                     \
									typename LinCopy::Src(input_offset, input_stride, src.data),                       \
									dst_inner_dim_size);                                                               \
                                                                                                                       \
		for (int j = 0; j < idx; ++j) {                                                                                \
			if (++it[j].count < it[j].size) {                                                                          \
				input_offset += it[j].input_stride;                                                                    \
				output_offset += it[j].output_stride;                                                                  \
				break;                                                                                                 \
			}                                                                                                          \
			it[j].count = 0;                                                                                           \
			input_offset -= it[j].input_span;                                                                          \
			output_offset -= it[j].output_span;                                                                        \
		}                                                                                                              \
	}                                                                                                                  \
	return num_copied;

		if (input_stride == 1 && output_stride == 1) {
			COPY_INNER_DIM(LinCopy::Kind::Linear);
		} else if (input_stride == 1 && output_stride != 1) {
			COPY_INNER_DIM(LinCopy::Kind::Scatter);
		} else if (input_stride == 0 && output_stride == 1) {
			COPY_INNER_DIM(LinCopy::Kind::FillLinear);
		} else if (input_stride == 0 && output_stride != 1) {
			COPY_INNER_DIM(LinCopy::Kind::FillScatter);
		} else if (output_stride == 1) {
			COPY_INNER_DIM(LinCopy::Kind::Gather);
		} else {
			COPY_INNER_DIM(LinCopy::Kind::Random);
		}

#undef COPY_INNER_DIM
	}

	// Copy from `src` to `dst` with an identity src->dst dimension map. Returns
	// the number of copied elements.
	static EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE IndexType Copy(const Dst& dst, const Src& src)
	{
		DimensionsMap dst_to_src_map;
		for (int i = 0; i < NumDims; ++i)
			dst_to_src_map[i] = i;
		return Copy(dst, src, dst_to_src_map);
	}

  private:
	struct BlockIteratorState
	{
		BlockIteratorState()
			: size(0)
			, count(0)
			, input_stride(0)
			, output_stride(0)
			, input_span(0)
			, output_span(0)
		{
		}

		IndexType size;
		IndexType count;
		IndexType input_stride;
		IndexType output_stride;
		IndexType input_span;
		IndexType output_span;
	};

	// Compute how many inner dimensions it's allowed to squeeze when doing IO
	// between two tensor blocks. It's safe to squeeze inner dimensions, only
	// if they are not reordered.
	static int NumSqueezableInnerDims(const DimensionsMap& dim_map)
	{
		int num_squeezable_dims = 0;
		for (int i = 0; i < NumDims; ++i) {
			const int dim = IsColMajor ? i : NumDims - i - 1;
			if (dim_map[dim] != dim)
				break;
			num_squeezable_dims++;
		}
		return num_squeezable_dims;
	}
};

// -------------------------------------------------------------------------- //
// TensorBlockAssignment assigns a block expression of type `TensorBlockExpr` to
// a Tensor block defined by `desc`, backed by a memory buffer at `target`.
//
// Currently there is no way to write from a Tensor expression to a block of
// memory, if dimensions are reordered. If you need to do that, you should
// materialize a Tensor block expression into a memory buffer, and then use
// TensorBlockIO to copy data between two memory buffers with a custom
// `target->src` dimension map (see definition above).
//
// Also currently the innermost dimension of `target` must have a stride '1'
// (contiguous in memory). This restriction could be lifted with a `pscatter`,
// but in practice it's never needed, and there is a similar TensorBlockIO
// workaround for that.
//
// TODO(ezhulenev): TensorBlockAssignment is a special case of TensorBlockIO
// where `src` is a tensor expression. Explore if it is possible to rewrite IO
// to use expressions instead of pointers, and after that TensorBlockAssignment
// will become an alias to IO.
template<typename Scalar, int NumDims, typename TensorBlockExpr, typename IndexType = Eigen::Index>
class TensorBlockAssignment
{
	// We will use coeff/packet path to evaluate block expressions.
	typedef TensorEvaluator<const TensorBlockExpr, DefaultDevice> TensorBlockEvaluator;

	typedef DSizes<IndexType, NumDims> Dimensions;

	enum
	{
		Vectorizable = packet_traits<Scalar>::Vectorizable,
		PacketSize = packet_traits<Scalar>::size
	};

	template<bool Vectorizable, typename Evaluator>
	struct InnerDimAssign
	{
		EIGEN_ALWAYS_INLINE static void Run(Scalar* target,
											IndexType count,
											const Evaluator& eval,
											IndexType eval_offset)
		{
			for (IndexType i = 0; i < count; ++i) {
				target[i] = eval.coeff(eval_offset + i);
			}
		}
	};

	template<typename Evaluator>
	struct InnerDimAssign<true, Evaluator>
	{
		EIGEN_ALWAYS_INLINE static void Run(Scalar* target,
											IndexType count,
											const Evaluator& eval,
											IndexType eval_offset)
		{
			typedef typename packet_traits<Scalar>::type Packet;

			const IndexType unrolled_size = count - 4 * PacketSize;
			const IndexType vectorized_size = count - PacketSize;
			IndexType i = 0;

			for (; i <= unrolled_size; i += 4 * PacketSize) {
				for (int j = 0; j < 4; ++j) {
					const IndexType idx = eval_offset + i + j * PacketSize;
					Packet p = eval.template packet<Unaligned>(idx);
					pstoreu<Scalar>(target + i + j * PacketSize, p);
				}
			}

			for (; i <= vectorized_size; i += PacketSize) {
				Packet p = eval.template packet<Unaligned>(eval_offset + i);
				pstoreu<Scalar>(target + i, p);
			}

			for (; i < count; ++i) {
				target[i] = eval.coeff(eval_offset + i);
			}
		}
	};

  public:
	struct Target
	{
		Target(const Dimensions& target_dims,
			   const Dimensions& target_strides,
			   Scalar* target_data,
			   IndexType target_offset = 0)
			: dims(target_dims)
			, strides(target_strides)
			, data(target_data)
			, offset(target_offset)
		{
		}

		Dimensions dims;
		Dimensions strides;
		Scalar* data;
		IndexType offset;
	};

	static Target target(const Dimensions& target_dims,
						 const Dimensions& target_strides,
						 Scalar* target_data,
						 IndexType target_offset = 0)
	{
		return Target(target_dims, target_strides, target_data, target_offset);
	}

	template<typename TargetDimsIndexType, typename TargetStridesIndexType>
	static Target target(const DSizes<TargetDimsIndexType, NumDims>& target_dims,
						 const DSizes<TargetStridesIndexType, NumDims>& target_strides,
						 Scalar* target_data,
						 IndexType target_offset = 0)
	{
		// DSizes constructor will do index type promotion if it's safe.
		return Target(Dimensions(target_dims), Dimensions(target_strides), target_data, target_offset);
	}

	static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void Run(const Target& target, const TensorBlockExpr& expr)
	{
		// Prepare evaluator for block expression.
		DefaultDevice default_device;
		TensorBlockEvaluator eval(expr, default_device);

		// Tensor block expression dimension should match destination dimensions.
		eigen_assert(dimensions_match(target.dims, eval.dimensions()));

		static const int Layout = TensorBlockEvaluator::Layout;
		static const bool is_col_major = Layout == ColMajor;

		// Initialize output inner dimension size based on a layout.
		const IndexType output_size = NumDims == 0 ? 1 : target.dims.TotalSize();
		const int inner_dim_idx = is_col_major ? 0 : NumDims - 1;
		IndexType output_inner_dim_size = target.dims[inner_dim_idx];

		// Target inner dimension stride must be '1'.
		eigen_assert(target.strides[inner_dim_idx] == 1);

		// Squeeze multiple inner dims into one if they are contiguous in `target`.
		IndexType num_squeezed_dims = 0;
		for (Index i = 1; i < NumDims; ++i) {
			const Index dim = is_col_major ? i : NumDims - i - 1;
			const IndexType target_stride = target.strides[dim];

			if (output_inner_dim_size == target_stride) {
				output_inner_dim_size *= target.dims[dim];
				num_squeezed_dims++;
			} else {
				break;
			}
		}

		// Initialize output block iterator state. Dimension in this array are
		// always in inner_most -> outer_most order (col major layout).
		array<BlockIteratorState, NumDims> it;

		int idx = 0; // currently initialized iterator state index
		for (Index i = num_squeezed_dims; i < NumDims - 1; ++i) {
			const Index dim = is_col_major ? i + 1 : NumDims - i - 2;

			it[idx].count = 0;
			it[idx].size = target.dims[dim];
			it[idx].output_stride = target.strides[dim];
			it[idx].output_span = it[idx].output_stride * (it[idx].size - 1);
			idx++;
		}

		// We read block expression from the beginning, and start writing data to
		// `target` at given offset.
		IndexType input_offset = 0;
		IndexType output_offset = target.offset;

		// Iterate copying data from `eval` to `target`.
		for (IndexType i = 0; i < output_size; i += output_inner_dim_size) {
			// Assign to `target` at current offset.
			InnerDimAssign<Vectorizable && TensorBlockEvaluator::PacketAccess, TensorBlockEvaluator>::Run(
				target.data + output_offset, output_inner_dim_size, eval, input_offset);

			// Move input offset forward by the number of assigned coefficients.
			input_offset += output_inner_dim_size;

			// Update index.
			for (int j = 0; j < idx; ++j) {
				if (++it[j].count < it[j].size) {
					output_offset += it[j].output_stride;
					break;
				}
				it[j].count = 0;
				output_offset -= it[j].output_span;
			}
		}
	}

  private:
	struct BlockIteratorState
	{
		BlockIteratorState()
			: count(0)
			, size(0)
			, output_stride(0)
			, output_span(0)
		{
		}

		IndexType count;
		IndexType size;
		IndexType output_stride;
		IndexType output_span;
	};
};

// -------------------------------------------------------------------------- //

} // namespace internal
} // namespace Eigen

#endif // EIGEN_CXX11_TENSOR_TENSOR_BLOCK_H
